/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of UHM (Unassembled HyperMatrix) sparse direct solver.
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#include "uhm/common.hxx"
#include "uhm/const.hxx"
#include "uhm/util.hxx"

#include "uhm/object.hxx"
#include "uhm/mesh/node.hxx"
#include "uhm/mesh/element.hxx"

#include "uhm/matrix/uhm/matrix.hxx"
#include "uhm/matrix/uhm/helper.hxx"

namespace uhm {
  // --------------------------------------------------------------
  // ** Helper

  // --------------------------------------------------------------
  // ** Definition
  Helper_::Helper_() {  }
  Helper_::Helper_(Element p, Element c) {
    assert( element_valid(p)       && element_valid(c) &&
	    p->is_matrix_created() && c->is_matrix_created() );

    this->cookie = UHM_HELPER_COOKIE;
    this->p      = p;    
    this->c      = c;
  }
  Helper_::~Helper_() { }

  // --------------------------------------------------------------
  void Helper_::set_mapper() {
    Mapper_ q;

    std::map< Node, int > factor, schur;

    // ** dump vector into map
    for (int i=0;i<this->p->factor.size();++i) 
      factor.insert( this->p->factor.at(i) );
    for (int i=0;i<this->p->schur.size();++i)
      schur.insert( this->p->schur.at(i) );

    std::vector<Mapper_> bijection;
    bijection.reserve(this->c->schur.size());

    // ** visit child schur nodes
    std::map< Node, int>::iterator pit;
    std::vector< std::pair<Node, int> >::iterator cit;
    for (cit=this->c->schur.begin();cit!=this->c->schur.end();++cit) {
      q.offs_c = cit->second;
      
      // find the node living in parent element
      pit = factor.find(cit->first);
      
      if (pit == factor.end()) {
	pit = schur.find(cit->first);
	assert(pit != schur.end());
	q.fs_p = 1; 
      } else {
	q.fs_p = 0; 
      }
      
      // set offset and dof
      q.offs_p = pit->second;
      q.n_dof  = cit->first->get_n_dof();
      
      // push into bijection
      bijection.push_back(q);
    }


    // condense
    this->mapper.clear();

    if (!bijection.size()) return;
    
    q = bijection.at(0);
    for (int i=0;i<bijection.size();++i) {
      if (i) {
	int flag = (bijection.at(i).fs_p   - bijection.at(i-1).fs_p);	  
	int diff = (bijection.at(i).offs_p - bijection.at(i-1).offs_p);
	
	if (!flag && (diff == bijection.at(i-1).n_dof)) {
	  this->mapper.back().n_dof += bijection.at(i).n_dof;
	} else {
	  this->mapper.push_back( bijection.at(i) );
	}
      } else {
	this->mapper.push_back( bijection.at(i) );
      }
    }
  }

  void Helper_::merge_A()     { _merge_A(); }
  void Helper_::branch_ABR()  { _branch_ABR(); }

  void Helper_::merge_rhs_x() { _merge_rhs(0,   false); }
  void Helper_::merge_rhs_b() { _merge_rhs(10,  false); }
  void Helper_::merge_rhs_r() { _merge_rhs(100, true); }

  void Helper_::branch_rhs_x() { _branch_rhs(0); }
  void Helper_::branch_rhs_b() { _branch_rhs(10); }
  void Helper_::branch_rhs_r() { _branch_rhs(100); }

  bool Helper_::disp() { 
    this->disp(stdout);
    return true;
  }
  bool Helper_::disp(FILE *stream) {
    fprintf(stream, "- Helper -\n");
    fprintf(stream, "  parent [ %d ], child [ %d ] \n",
	    this->p->get_id(), this->c->get_id());
    
    std::vector<Mapper_>::iterator it;
    for (it=this->mapper.begin();it<this->mapper.end();++it) 
      fprintf(stream, 
	      "  side :: [ %d ], offset :: ( %d , %d ), n_dof :: < %d >\n",
	      it->fs_p, it->offs_p, it->offs_c, it->n_dof);
    
    return true;
  }

  // --------------------------------------------------------------
  // ** Protected 
  void Helper_::_merge_A() {
    if (!this->mapper.size()) return;

    Matrix parent = this->p->get_matrix();
    Matrix child  = this->c->get_matrix();
    int is_erase  = false;

    int n         = this->mapper.size();

#ifdef UHM_LIGHT_BLAS_ONE_ENABLE
    int datatype = parent->get_datatype(), cs[5];
    std::pair<double*,int> buf[5];
    buf[0] = child->get_buffer(UHM_ABR);
    for (int i=UHM_ATL;i<UHM_P;++i) 
      buf[i] = parent->get_buffer(i);

#endif

#pragma unroll(UHM_UNROLL_N)
    for (int j = 0 ;j < n; ++j) {

#pragma unroll(UHM_UNROLL_N)
      for (int i = 0 ;i < n; ++i) {

        int mat     = _get_A(this->mapper.at(i).fs_p,
                             this->mapper.at(j).fs_p);

        if (parent->is_symmetry() && mat == UHM_ATR) continue;

        int ioffs_c = this->mapper.at(i).offs_c; int joffs_c = this->mapper.at(j).offs_c;
        int ioffs_p = this->mapper.at(i).offs_p; int joffs_p = this->mapper.at(j).offs_p;
        int in_dof  = this->mapper.at(i).n_dof;  int jn_dof  = this->mapper.at(j).n_dof;

#ifdef UHM_LIGHT_BLAS_ONE_ENABLE
        merge(datatype,
              buf[0].first,   ioffs_c, joffs_c, buf[0].second,
              buf[mat].first, ioffs_p, joffs_p, buf[mat].second,
              in_dof, jn_dof,
              is_erase);
#else
        parent->merge( child,
                       UHM_ABR,
                       ioffs_c, joffs_c,
                       mat,
                       ioffs_p, joffs_p,
                       in_dof, jn_dof,
                       is_erase);
#endif
      }
    }
  }
  void Helper_::_branch_ABR() {
    if (!this->mapper.size()) return;

    Matrix parent = this->p->get_matrix();
    Matrix child  = this->c->get_matrix();
    int is_erase  = false;

    int n         = this->mapper.size();

#pragma unroll(UHM_UNROLL_N)
    for (int j = 0 ;j < n; ++j) {

#pragma unroll(UHM_UNROLL_N)
      for (int i = 0 ;i < n; ++i) {

        int mat     = _get_A(this->mapper.at(i).fs_p, this->mapper.at(j).fs_p);
        int ioffs_c = this->mapper.at(i).offs_c; int joffs_c = this->mapper.at(j).offs_c;
        int ioffs_p = this->mapper.at(i).offs_p; int joffs_p = this->mapper.at(j).offs_p;
        int in_dof  = this->mapper.at(i).n_dof;  int jn_dof  = this->mapper.at(j).n_dof;

        child->copy( parent,
                     mat,
                     ioffs_p, joffs_p,
                     UHM_ABR,
                     ioffs_c, joffs_c,
                     in_dof, jn_dof,
                     is_erase);
      }
    }
  }
  void Helper_::_merge_rhs(int kind, int is_pivot_applied) {
    if (!this->mapper.size()) return;

    // for checking routine it is necessary to apply pivot
    // at this level -- tricky...
    if (is_pivot_applied &&
        this->p->get_matrix()->is_created( _get_rhs(kind) ) &&
        this->p->get_matrix()->is_buffer( _get_rhs(kind) ) ){
      this->p->get_matrix()->backup   ( _get_rhs(kind) );
      this->p->get_matrix()->set_zero ( _get_rhs(kind) );
    }

    Matrix parent   = this->p->get_matrix();
    Matrix child    = this->c->get_matrix();
    int    is_erase = true;
    int    n_rhs    = p->get_matrix()->get_n_rhs();

    int n           = this->mapper.size();

#pragma unroll(UHM_UNROLL_N)
    for (int j = 0; j < n ; ++j) {
      int mat_src = _get_rhs(1+kind);
      int mat_tgt = _get_rhs(this->mapper.at(j).fs_p+kind);

      int joffs_c = this->mapper.at(j).offs_c;
      int joffs_p = this->mapper.at(j).offs_p;

      int jn_dof  = this->mapper.at(j).n_dof;

      parent->merge( child, 
		     mat_src,
                     joffs_c, 0,
                     mat_tgt,
                     joffs_p, 0,
                     jn_dof, n_rhs, 
		     is_erase);

      // after merge set child bb zero
    }

    if (is_pivot_applied &&
        this->p->get_matrix()->is_created( _get_rhs(kind) ) &&
        this->p->get_matrix()->is_buffer( _get_rhs(kind) ) ){
      int is_merge = true;

      this->p->get_matrix()->apply_pivots( _get_rhs(kind) );
      this->p->get_matrix()->restore( _get_rhs(kind), is_merge );
    }
  }
  void Helper_::_branch_rhs(int kind) {
    if (!this->mapper.size()) return;

    Matrix parent   = this->p->get_matrix();
    Matrix child    = this->c->get_matrix();
    int    is_erase = false;
    int    n_rhs    = p->get_matrix()->get_n_rhs();

    int n           = this->mapper.size();

#pragma unroll(UHM_UNROLL_N)
    for (int j = 0; j < n ; ++j) {

      int mat_src = _get_rhs(this->mapper.at(j).fs_p+kind);
      int mat_tgt = _get_rhs(1+kind);

      int joffs_c = this->mapper.at(j).offs_c;
      int joffs_p = this->mapper.at(j).offs_p;

      int jn_dof  = this->mapper.at(j).n_dof;

      child->copy( parent,
                   mat_src,
                   joffs_p, 0,
                   mat_tgt,
                   joffs_c, 0,
                   jn_dof, n_rhs, 
		   is_erase );
    }
  }

  int Helper_::_get_rhs(int col) {
    switch (col) {
    case 0:   return UHM_XT;
    case 1:   return UHM_XB;
    case 10:  return UHM_BT;
    case 11:  return UHM_BB;
    case 100: return UHM_RT;
    case 101: return UHM_RB;
    }
    return 0;
  }
  int Helper_::_get_A(int col, int row) {
    switch (col*10+row) {
    case 0:  return UHM_ATL;
    case 1:  return UHM_ATR;
    case 10: return UHM_ABL;
    case 11: return UHM_ABR;
    }
    return 0;
  }
}
